import sys
from PyQt5.QtWidgets import QApplication , QMainWindow
from FirstMainWin import *
from img2hex2 import *
from hex2img2 import *
from sm4Decode import *
from LSBde import * 
# import os

class MyMainWindow(QMainWindow, Ui_MainWindow):    
    global list1,list2
    list1 = []
    list2 = []
    def __init__(self,parent=None):        
        super(MyMainWindow, self).__init__(parent)     
        self.setupUi(self)
        # self.open_file1()
        self.pushButton.clicked.connect(self.open_file1)
        self.pushButton_2.clicked.connect(self.input_width)
        self.pushButton_3.clicked.connect(self.input_height)
        self.pushButton_5.clicked.connect(self.input_key)
        self.pushButton_4.clicked.connect(self.jiemi)


    def open_file1(self):
        fileName = QtWidgets.QFileDialog.getOpenFileName(self,"选取待解密图片","e:/")
        self.listWidget.addItem('您所选择的待解密图片为'+fileName[0])
        list1.append(fileName[0])
    def input_width(self):
        width = QtWidgets.QInputDialog.getText(self,"请输入","请输入水印图片的宽(数字水印解密密钥哦~)")
        self.listWidget_2.addItem('已经输入水印图片的宽')
        list2.append(int(width[0]))
    def input_height(self):
        height = QtWidgets.QInputDialog.getText(self,"请输入","请输入水印图片的高(数字水印解密密钥哦~)")
        self.listWidget_3.addItem('已经输入水印图片的高')
        list2.append(int(height[0]))
    def input_key(self):
        key = QtWidgets.QInputDialog.getText(self,"请输入","请输入SM4密钥")
        self.listWidget_4.addItem('已经输入SM4密钥')
        list1.append(key[0])

    def jiemi(self):
        filepath1 = 'e:\\mms\\cryptoout.txt'
        filepath2 = 'e:\\mms\\plaintstout.txt'
        imgpath1 = 'e:\\mms\\还原加密水印图片.bmp'
        img1 = Image.open(list1[0])
        rgb_im1 = img1.convert('RGB')
        length = list2[0] * list2[1] * 24
        wt = deEncry(rgb_im1, length)
        showImage(wt, list2[0], list2[1], imgpath1)
        img2 = Image.open(imgpath1)
        # 将像素值转为文本
        rgb_img = img2.convert('RGB')
        hexcode = getCode(rgb_img)
        # print(len(hexcode))
        with open(filepath1, 'w') as file_object:
            file_object.write(hexcode)
        key = bytes.fromhex(list1[1])
        sm4 = SM4Cipher(key)
        with open(filepath1) as fp1:
                crypt = fp1.read()
        fp2 = open(filepath2,'a+')
        m = list2[0]*list2[1]*6
        a = sm4.decrypt(bytes.fromhex(crypt[len(crypt)-32:])).hex()
        for i in range(0, m-32, 32):
            fp2.write(sm4.decrypt(bytes.fromhex(crypt[i:i+32])).hex())
        fp2.write(a[32-(m%32):])
        fp2.close()
        with open(filepath2) as fp:
            a = fp.read()
        path1 = list1[0] + '/../还原水印图片.bmp'
        showImage1(a, list2[0],list2[1],path1)
        self.listWidget_5.addItem('解密成功!已输出到同级文件夹!')
        print('解密成功')
    

if __name__ == '__main__':     
    app = QApplication(sys.argv)
    myWin = MyMainWindow()
    myWin.show()
    sys.exit(app.exec_())